import { promises as fs } from "fs";
import { resolve } from "path";
import { ChatAgentConfig } from "../../agent/chatAgent";
import { CodeAbstractorAgent, CodeFolderAbstractorAgent } from "../../agent/codeAgent";
import { Language, RoleType } from "../../enum";
import { ReceivingChatMessage } from "../../message/chatMessage";
import { getPromptFromKey, promptTemplateDict } from "../../prompt";
import { personConfigs } from "../../config";
import { ALLOW_FILE_TYPES } from "./config";
import {
    ProjectProcessor,
    ProjectItem,
    ItemTypeEnum,
    ProjectItemFolder,
    ProjectItemFile
} from "./projectProcessor";
import { ProcessRecorder } from "./processRecorder";

export interface ProjectParsingConfig {
    /* 传递给助手代理的额外参数，默认为 null */
    agentConfig?: ChatAgentConfig;
    /* 公共记忆，默认为 null */
    memory?: any;
    /* 记忆深度，默认为 1 */
    memoryDepth?: number;
    /* 语言，默认 Chinese */
    language?: Language;
    /** 最大文件大小，默认为 10MB */
    maxFileSize?: number;
    /** 文件编码，默认为 utf-8 */
    fileEncoding?: FileEncodingEnum;
    /** 项目名称 */
    projectName?: string;
    /** 项目介绍 */
    projectDescription?: string;
}

enum FileEncodingEnum {
    UTF8 = "utf-8"
}

export class ProjectParsing {
    public agentConfig?: ChatAgentConfig;
    public memory?: any;
    public language?: Language;
    public maxFileSize: number;
    public fileEncoding: FileEncodingEnum;
    public projectName: string;
    public projectDescription: string;
    public projectPath: string = "";

    /** 项目摘要 */
    public abstracts: ProjectItem[] = [];

    /** Agent 字典 */
    private agentDict = {
        codeAbstractorAgent: CodeAbstractorAgent,
        codeFolderAbstractorAgent: CodeFolderAbstractorAgent
    };
    /** Agent 缓存字典 */
    private agentTempDict: {
        codeAbstractorAgent?: CodeAbstractorAgent;
        codeFolderAbstractorAgent?: CodeFolderAbstractorAgent;
    } = {};

    constructor(config: ProjectParsingConfig) {
        this.agentConfig = config.agentConfig;
        this.memory = config.memory || null;
        this.language = config.language ?? personConfigs.defaultLanguage;
        this.maxFileSize = config.maxFileSize ?? 10 * 1024 * 1024;
        this.fileEncoding = config.fileEncoding ?? FileEncodingEnum.UTF8;
        this.projectName = config.projectName ?? "";
        this.projectDescription = config.projectDescription ?? "";
    }

    async initChat() {
        const codeAbstractorAgent = this.agentTempDict["codeAbstractorAgent"];
        const codeFolderAbstractorAgent = this.agentTempDict["codeFolderAbstractorAgent"];
        if (codeAbstractorAgent && codeFolderAbstractorAgent) {
            codeAbstractorAgent.reset();
            codeFolderAbstractorAgent.reset();
        } else {
            this.agentTempDict["codeAbstractorAgent"] = new this.agentDict["codeAbstractorAgent"]({
                ...this.agentConfig,
                language: this.language,
                memory: this.memory,
                projectName: this.projectName,
                projectDescription: this.projectDescription
            });
            this.agentTempDict["codeFolderAbstractorAgent"] = new this.agentDict[
                "codeFolderAbstractorAgent"
            ]({
                ...this.agentConfig,
                language: this.language,
                memory: this.memory,
                projectName: this.projectName,
                projectDescription: this.projectDescription
            });
        }
        return {
            codeAbstractorAgent: this.agentTempDict["codeAbstractorAgent"]!,
            codeFolderAbstractorAgent: this.agentTempDict["codeFolderAbstractorAgent"]!
        };
    }

    /**
     * 解析项目
     * @param projectPath 项目路径
     * @param update 是否更新
     * @returns 项目摘要
     *
     * @description
     * 传入路径与缓存路径有以下几种情况:
     * - 第一次解析，建立摘要
     * - 不相等、且传入路径为子路径、更新摘要
     * - 不相等、且传入路径为子路径、不更新摘要
     * - 不相等、且传入路径不为子路径、清空摘要，重新解析
     * - 相等、更新摘要
     * - 相等、不更新摘要
     */
    async step(path: string, update: boolean = false) {
        const projectPath = resolve(path);

        if (!this.projectPath) {
            // 第一次解析，建立摘要
            return await this.updateAbstracts(projectPath);
        }
        if (this.projectPath !== projectPath) {
            if (this.isSubPath(projectPath, this.projectPath)) {
                if (update) {
                    // 不相等、且传入路径为子路径、更新摘要
                    return await this.updateAbstracts(projectPath);
                } else {
                    // 不相等、且传入路径为子路径、不更新摘要
                    return await this.getSubAbstracts(projectPath);
                }
            } else {
                // 不相等、且传入路径不为子路径、清空摘要，重新解析
                this.abstracts = [];
                this.projectPath = projectPath;
                return await this.updateAbstracts(projectPath);
            }
        }
        if (update) {
            // 相等、更新摘要
            return await this.updateAbstracts(projectPath);
        } else {
            // 相等、不更新摘要
            return this.abstracts;
        }
    }

    private async getSubAbstracts(path: string) {
        let flag = true;
        const items = await ProjectProcessor.traverseDirectory(path, (item) => {
            if (item.absolutePath === path) {
                flag = false;
            }
            return item;
        });
        if (flag) {
            throw new Error(`路径 ${path} 为空`);
        }
        return items[0];
    }

    private async updateAbstracts(path: string) {
        this.abstracts = await this.analyzeProject(path, true);
        return this.abstracts;
    }

    private isSubPath(subPath: string, mainPath: string) {
        // 转换为绝对路径
        const absSubPath = resolve(subPath);
        const absMainPath = resolve(mainPath);

        const normalizedAbsSubPath = absSubPath.toLowerCase();
        const normalizedAbsMainPath = absMainPath.toLowerCase();

        return normalizedAbsSubPath.startsWith(normalizedAbsMainPath);
    }

    private async analyzeProject(path: string, update: boolean): Promise<ProjectItem[]> {
        const { codeAbstractorAgent, codeFolderAbstractorAgent } = await this.initChat();
        const processRecorder = new ProcessRecorder();

        processRecorder.logState("初始化完成");

        // 验证文件
        const processedFiles1 = await ProjectProcessor.traverseDirectory(path, async (item) => {
            const result = await this.validateProjectFileAndFolder(item);
            if (!result) {
                const errorMessage = `文件 ${item.absolutePath} 未通过校验`;
                item.error = { message: errorMessage, skip: true };
            }
            processRecorder.logAddItem(item);
            return item;
        });

        processRecorder.logState("项目文件校验完成");

        // 处理文件
        const processedFiles2 = await this.analyzeProjectFiles(
            processedFiles1,
            codeAbstractorAgent,
            processRecorder,
            update
        );

        // 处理文件夹
        const processedFiles3 = await this.analyzeProjectFolders(
            processedFiles2,
            codeFolderAbstractorAgent,
            processRecorder,
            update
        );

        processRecorder.logState("项目摘要生成完成");

        return processedFiles3;
    }

    private async analyzeProjectFiles(
        lastProcessedFiles: ProjectItem[],
        codeAbstractorAgent: CodeAbstractorAgent,
        processRecorder: ProcessRecorder,
        update: boolean
    ) {
        // 处理文件
        const nextProcessedFiles = await ProjectProcessor.traverseDirectory(
            lastProcessedFiles,
            async (item: ProjectItem) => {
                if (item.error) {
                    return item;
                }
                if (item.type === ItemTypeEnum.File) {
                    const oldAbstract = item.abstract;
                    try {
                        const abstract = await this.analyzeProjectFile(
                            item,
                            codeAbstractorAgent,
                            update
                        );
                        const newItem = {
                            ...item,
                            abstract,
                            oldAbstract
                        };
                        processRecorder.logProcessItem(newItem);
                        return newItem;
                    } catch (error) {
                        const errorMessage = `文件 ${item.absolutePath} 摘要生成失败，原因：${(error as Error).message}`;
                        item.error = { message: errorMessage };
                        processRecorder.logProcessItem(item);
                        return item;
                    }
                } else {
                    return item;
                }
            }
        );
        return nextProcessedFiles;
    }

    private async analyzeProjectFolders(
        lastProcessedFiles: ProjectItem[],
        codeFolderAbstractorAgent: CodeFolderAbstractorAgent,
        processRecorder: ProcessRecorder,
        update: boolean
    ) {
        // 处理文件夹
        const nextProcessedFiles = await ProjectProcessor.traverseDirectory(
            lastProcessedFiles,
            async (item: ProjectItem) => {
                if (item.error) {
                    return item;
                }
                if (item.type === ItemTypeEnum.File) {
                    return item;
                } else {
                    const oldAbstract = item.abstract;
                    try {
                        const abstract = await this.analyzeProjectFolder(
                            item,
                            codeFolderAbstractorAgent,
                            update
                        );
                        const newItem = {
                            ...item,
                            abstract,
                            oldAbstract
                        };
                        processRecorder.logProcessItem(newItem);
                        return newItem;
                    } catch (error) {
                        const errorMessage = `文件夹 ${item.absolutePath} 摘要生成失败，原因：${(error as Error).message}`;
                        item.error = { message: errorMessage };
                        processRecorder.logProcessItem(item);
                        return item;
                    }
                }
            }
        );
        return nextProcessedFiles;
    }

    private async validateProjectFileAndFolder(item: ProjectItem): Promise<boolean> {
        const path = item.absolutePath;
        if (!(await this.validateFile(path))) {
            return false;
        }
        return true;
    }

    private async analyzeProjectFile(
        item: ProjectItemFile,
        agent: CodeAbstractorAgent,
        update: boolean
    ) {
        if (!update) {
            if (item.abstract) return item.abstract;
        }
        const path = item.absolutePath;
        const code = await fs.readFile(path, this.fileEncoding);
        const response = await agent.step(
            new ReceivingChatMessage({
                roleName: "ProjectParsing",
                roleType: RoleType.User,
                content: this.handelFileTemplate(code, item)
            })
        );
        if (response.terminated) {
            throw new Error(`对话已结束，请重新初始化对话`);
        }
        return this.handelContent(response.messages.slice(-1)[0].content, item);
    }

    private async analyzeProjectFolder(
        item: ProjectItemFolder,
        agent: CodeFolderAbstractorAgent,
        update: boolean
    ) {
        if (!update) {
            if (item.abstract) return item.abstract;
        }
        const content = item.subItems
            .map((item) => {
                return item.abstract;
            })
            .join("\n ");
        const response = await agent.step(
            new ReceivingChatMessage({
                roleName: "ProjectParsing",
                roleType: RoleType.User,
                content: this.handelFolderTemplate(content, item)
            })
        );
        if (response.terminated) {
            throw new Error(`对话已结束，请重新初始化对话`);
        }
        return this.handelContent(response.messages.splice(-1)[0].content, item);
    }

    private handelFileTemplate(fileContent: string, metaData: ProjectItem) {
        return getPromptFromKey(
            promptTemplateDict.projectParsingTemplate.fileAdditional,
            this.language
        ).format({
            code: fileContent,
            projectName: this.projectName,
            path: metaData.absolutePath
        });
    }

    private handelFolderTemplate(content: string, metaData: ProjectItem) {
        return getPromptFromKey(
            promptTemplateDict.projectParsingTemplate.folderAdditional,
            this.language
        ).format({
            content: content,
            projectName: this.projectName,
            path: metaData.absolutePath
        });
    }

    private handelContent(content: string, metaData: ProjectItem) {
        return getPromptFromKey(
            promptTemplateDict.projectParsingTemplate.itemMetaData,
            this.language
        ).format({
            projectName: this.projectName,
            path: metaData.absolutePath,
            abstract: content
        });
    }

    private async validateFile(path: string): Promise<boolean> {
        const stats = await fs.stat(path);
        // 校验是否为文件
        if (!stats.isFile()) return true;

        const fileSize = stats.size;
        // 校验文件大小
        if (fileSize > this.maxFileSize) {
            return false;
        }
        // 校验文件类型
        const ext = path.split(".").pop();
        if (!ext) {
            return false;
        }
        if (ALLOW_FILE_TYPES.findIndex((item) => item === ext.toLowerCase()) === -1) {
            return false;
        }
        return true;
    }
}
